import os.path
import random

raw_file = 'mathbench_v1/college/single_choice_cn.jsonl'

with open(raw_file) as f:
    ls = f.readlines()
idxs = random.sample(range(len(ls)), 100)
print(sorted(idxs))
train_ls = []
eval_ls = []
for i in range(len(ls)):
    if i not in idxs:
        eval_ls.append(ls[i])
    else:
        train_ls.append(ls[i])
with open(os.path.join(os.path.dirname(raw_file), 'train.jsonl'), 'w') as f:
    f.writelines(train_ls)
with open(os.path.join(os.path.dirname(raw_file), 'eval.jsonl'), 'w') as f:
    f.writelines(eval_ls)